package com.ly.pic.service;
import lombok.extern.slf4j.Slf4j;
import org.datavec.image.loader.NativeImageLoader;
import org.deeplearning4j.nn.graph.ComputationGraph;
import org.deeplearning4j.nn.transferlearning.TransferLearningHelper;
import org.deeplearning4j.zoo.ZooModel;
import org.deeplearning4j.zoo.model.VGG19;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.dataset.DataSet;
import org.nd4j.linalg.dataset.api.preprocessor.DataNormalization;
import org.nd4j.linalg.dataset.api.preprocessor.VGG16ImagePreProcessor;
import org.nd4j.linalg.ops.transforms.Transforms;
import org.springframework.stereotype.Service;

import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.List;
import java.util.concurrent.CopyOnWriteArrayList;

@Slf4j
@Service  //自动注入spring容器会产生兼容问题，改为手动注入
public class FeatureExtractionAGG19Service {

    private ComputationGraph vgg19;
    private TransferLearningHelper transferLearningHelper;
    private NativeImageLoader loader;
    private DataNormalization scaler;
    public List<INDArray> featureList = new CopyOnWriteArrayList<>();
    public FeatureExtractionAGG19Service() throws IOException {
        // 指定本地模型文件的路径
        File modelFile = new File("models/vgg19_dl4j_inference.zip");

        // 检查模型文件是否存在
        if (modelFile.exists()) {
            vgg19 = ComputationGraph.load(modelFile, true);
        } else {
            ZooModel zooModel = VGG19.builder().build();
            vgg19 = (ComputationGraph) zooModel.initPretrained();
        }
        transferLearningHelper = new TransferLearningHelper(vgg19, "fc2");
        loader = new NativeImageLoader(224, 224, 3);
        scaler = new VGG16ImagePreProcessor();
    }

    public static void calculateTotalMemoryUsage(List<INDArray> featureList) {
        long totalBytes = 0;
        for (INDArray features : featureList) {
            long numElements = features.length();
            long bytesPerElement = features.data().getElementSize();
            long memoryUsage = numElements * bytesPerElement;
            totalBytes += memoryUsage;
        }
        double totalMemoryMB = totalBytes / (1024.0 * 1024.0);
        log.info("feature Total memory usage: {} MB", totalMemoryMB);
    }

    public INDArray extractFeatures(InputStream inputStream) throws IOException {
        INDArray image = loader.asMatrix(inputStream);
        try {
            scaler.transform(image);
            DataSet dataSet = new DataSet(image, null);
            return transferLearningHelper.featurize(dataSet).getFeatures();
        } finally {
            image.close();
        }
    }

    public boolean addFeatures(INDArray indArray){
        if (comparativeFeatureVector(indArray,0.999)) {
            return false;
        }
        return featureList.add(indArray);
    }

    public void setFeature(String imagePath) throws IOException {
        File file = new File(imagePath);
        try (INDArray image = loader.asMatrix(file)) {
            scaler.transform(image);
            DataSet dataSet = new DataSet(image, null);
            DataSet featurize = transferLearningHelper.featurize(dataSet);
            INDArray features = featurize.getFeatures();
            featureList.add(features);
            calculateTotalMemoryUsage(this.featureList);
        }finally {
            System.gc();
        }
    }

    public boolean comparativeFeatureVector(INDArray indArray, double similarity) {
        double sy = 0;
        for (INDArray array : featureList) {
            sy = Transforms.cosineSim(array, indArray);
            if (sy > similarity) {
                log.info("余弦相似度:{} 大于 输入相似度:{}",sy,similarity);
                return true;
            }
        }
        log.info("余弦相似度:{} 小于 输入相似度:{}",sy,similarity);
        return false;
    }



    public static void main(String[] args) throws IOException {
        long startTime = System.currentTimeMillis();
        FeatureExtractionAGG19Service featureExtractionAGG19Service = new FeatureExtractionAGG19Service();
        for (int i = 0; i < 100; i++) {
            INDArray features = featureExtractionAGG19Service.extractFeatures("C:\\Users\\ly\\Desktop\\resuouce\\pic-match\\004 - copy.png");
            INDArray features1 = featureExtractionAGG19Service.extractFeatures("C:\\Users\\ly\\Desktop\\resuouce\\pic-match\\004 - copy.png");
            long endTime = System.currentTimeMillis();
            if (features1.equals(features)) {
                System.out.println(true);
            }
            System.out.println("Time taken: " + (endTime - startTime) + " ms");
            System.out.println("Extracted Features: " + features);
        }

    }

    public INDArray extractFeatures(String imagePath) throws IOException {
        File file = new File(imagePath);
        try (InputStream inputStream = new FileInputStream(file)) {
            INDArray image = loader.asMatrix(inputStream);
            scaler.transform(image);
            DataSet dataSet = new DataSet(image, null);
            INDArray features = transferLearningHelper.featurize(dataSet).getFeatures();
            return features;
        }
    }

    public void addAllFeature(List<INDArray> List) {
        featureList.addAll(List);
    }
}
